import torch
import tokenizer
print(torch.cuda.is_available())

device = torch.device('cuda:0')
# 1. 加载模型
net = torch.load('model/net.pth', weights_only=False)
name_tk = tokenizer.get_name_tk()
cate_tk = tokenizer.get_cate_tk()
def predict(name):
    # 3. 利用名字tokenizer将名字转换为索引,注意要padding
    name_ids = name_tk.encode(name)
    name_ids = [0] * (19 - len(name_ids)) + name_ids
    net.eval()
    with torch.no_grad():
        name_ids_tensor = torch.LongTensor(name_ids).to(device) # (seq_len)
        y_pred = net(name_ids_tensor) #( label_num )
        # (seq_len)
    # 5. 取出top3的索引
        _, idxs = torch.topk(y_pred,k=3) # k=3意味着取前3个最大值所在的索引
    # 6. 将索引通过类别tokenizer转换为类别输出
        cates = cate_tk.decode(idxs.tolist())
    return cates

if __name__ == '__main__':
    c = predict('Rex')
    print(c)